# !/usr/bin/python3
# -*- coding: utf-8 -*-
# @Author : Lawrence
# @File : split_pdf.py
# @Project : worktool
# @Time : 2024/10/29 15:30
import os

import PyPDF2


def split_pdf(input_pdf, output_pdf1, output_pdf2, split_page):
    # 打开输入的PDF文件
    with open(input_pdf, 'rb') as file:
        reader = PyPDF2.PdfReader(file)

        # 创建两个输出PDF文件的写入对象
        writer1 = PyPDF2.PdfWriter()
        writer2 = PyPDF2.PdfWriter()

        # 遍历PDF的所有页面
        for page_num in range(len(reader.pages)):
            # page = reader.getPage(page_num)
            page = reader.pages[page_num]

            if page_num < split_page:
                # 将页面添加到第一个输出PDF文件
                writer1.add_page(page)
            else:
                # 将页面添加到第二个输出PDF文件
                writer2.add_page(page)

        # 将第一个输出PDF文件写入磁盘
        with open(output_pdf1, 'wb') as output_file1:
            writer1.write(output_file1)

        # 将第二个输出PDF文件写入磁盘
        with open(output_pdf2, 'wb') as output_file2:
            writer2.write(output_file2)


# 示例用法
input_pdf = 'data/GJB 438C-2021.pdf'  # 输入的PDF文件路径

output_pdf1 = os.path.dirname(input_pdf)+'/split1_'+os.path.basename(input_pdf)  # 输出的第一个PDF文件路径
output_pdf2 = os.path.dirname(input_pdf)+'/split2_'+os.path.basename(input_pdf)  # 输出的第二个PDF文件路径
split_page = 40  # 从第几页开始切分，例如第3页

split_pdf(input_pdf, output_pdf1, output_pdf2, split_page)